﻿#pragma once

#include <learn/alg/alg_approximator.h>
#include <learn/alg/alg_iterator.h>
#include <learn/alg/alg_projection.h>

#include <Dense>
#include <Sparse>

namespace xi
{

namespace alg
{

using namespace xi::math;
using namespace xi::geom;

template <std::size_t N> struct Alg_DBS_Approximator : Alg_Iterator<std::shared_ptr<Geom_BSplineSurface<N>>>
{
    // Dynamic Base Surface Approximator
    using Vec = VecT<N, double>; // P

    std::size_t uDegree{3}, vDegree{3};
    std::size_t uPoleNum{6}, vPoleNum{6};
    double wStretch{0.3 * 1e-6};
    double wBending{0.7 * 1e-6};
    std::vector<Vec> data;
    std::vector<Vec2> params;

  protected:
    std::shared_ptr<Geom_BSplineSurface<N>> m_surface;

    /**
     * @brief 初始参数化
     *
     */
    void init() override
    {
        if (params.empty())
            params = Alg_Parameter<N>::projection(data);
    }

    /**
     * @brief 迭代计算下一个曲面
     *
     * @return std::shared_ptr<Geom_BSplineSurface<N>>
     */
    std::shared_ptr<Geom_BSplineSurface<N>> next() override
    {
        auto n1 = uPoleNum - 1;
        auto n2 = vPoleNum - 1;
        auto m = data.size() - 1;

        auto uKnots = Alg_Knot<N>::equally(n1 + 1, uDegree);
        auto vKnots = Alg_Knot<N>::equally(n2 + 1, vDegree);

        // 最小二乘矩阵 A: (m + 1) x ((n1 + 1) x (n2 + 1))
        std::vector<Eigen::Triplet<double>> tripletsA;
        for (std::size_t i = 0; i < m + 1; i++)
        {
            Vec2 P = params[i];
            auto uSpan = find_span(uKnots.data(), uKnots.size(), uDegree, P.x());
            auto vSpan = find_span(vKnots.data(), vKnots.size(), vDegree, P.y());

            auto ubasis = basis_funcs(uKnots.data(), uDegree, uSpan, P.x());
            auto vbasis = basis_funcs(vKnots.data(), vDegree, vSpan, P.y());

            for (std::size_t j = 0; j < uDegree + 1; j++)
            {
                auto row = (uSpan - uDegree + j) * (n2 + 1);
                for (std::size_t k = 0; k < vDegree + 1; k++)
                {
                    auto col = row + (vSpan - vDegree + k);
                    tripletsA.emplace_back(i, col, ubasis[j] * vbasis[k]);
                }
            }
        }

        // 最小化弦长和导数的边界，显著快于能量项矩阵
        auto alpha = wStretch, beta = wBending, gamma = 1 - alpha - beta;
        std::vector<Eigen::Triplet<double>> tripletsB12;
        for (std::size_t i = 0; i < n1 + 1; i++)
        {
            int ii = i * (n2 + 1);

            // 最小化弦长的边界
            tripletsB12.emplace_back(ii, ii, alpha);
            tripletsB12.emplace_back(ii, ii + 1, -alpha);
            tripletsB12.emplace_back(ii + n2, ii + n2, alpha);
            tripletsB12.emplace_back(ii + n2, ii + n2 - 1, -alpha);

            // 最小化导数的边界
            tripletsB12.emplace_back(ii, ii, beta);
            tripletsB12.emplace_back(ii, ii + 1, -2 * beta);
            tripletsB12.emplace_back(ii, ii + 2, beta);

            tripletsB12.emplace_back(ii + 1, ii, -2 * beta);
            tripletsB12.emplace_back(ii + 1, ii + 1, 5 * beta);
            tripletsB12.emplace_back(ii + 1, ii + 2, -4 * beta);
            tripletsB12.emplace_back(ii + 1, ii + 3, beta);

            tripletsB12.emplace_back(ii + n2, ii + n2, beta);
            tripletsB12.emplace_back(ii + n2, ii + n2 - 1, -2 * beta);
            tripletsB12.emplace_back(ii + n2, ii + n2 - 2, beta);

            tripletsB12.emplace_back(ii + n2 - 1, ii + n2, -2 * beta);
            tripletsB12.emplace_back(ii + n2 - 1, ii + n2 - 1, 5 * beta);
            tripletsB12.emplace_back(ii + n2 - 1, ii + n2 - 2, -4 * beta);
            tripletsB12.emplace_back(ii + n2 - 1, ii + n2 - 3, beta);

            // 最小化弦长的中间部分
            for (std::size_t j = 1; j < n2; j++)
            {
                tripletsB12.emplace_back(ii + j, ii + j, 2 * alpha);
                tripletsB12.emplace_back(ii + j, ii + j - 1, -alpha);
                tripletsB12.emplace_back(ii + j, ii + j + 1, -alpha);
            }

            // 最小化导数的中间部分
            for (std::size_t j = 2; j < n2 - 1; j++)
            {
                tripletsB12.emplace_back(ii + j, ii + j, 6 * beta);
                tripletsB12.emplace_back(ii + j, ii + j - 1, -4 * beta);
                tripletsB12.emplace_back(ii + j, ii + j + 1, -4 * beta);
                tripletsB12.emplace_back(ii + j, ii + j - 2, beta);
                tripletsB12.emplace_back(ii + j, ii + j + 2, beta);
            }

            // 添加最小化弦长的中间部分
            if (i > 0 && i < n1)
            {
                for (std::size_t j = 0; j < n2 + 1; j++)
                {
                    tripletsB12.emplace_back(ii + j, ii + j, 2 * alpha);
                    tripletsB12.emplace_back(ii + j, ii + j + n2 + 1, -alpha);
                    tripletsB12.emplace_back(ii + j, ii + j - (n2 + 1), -alpha);
                }
            }

            // 添加最小化导数的中间部分
            if (i > 1 && i < n1 - 1)
            {
                for (std::size_t j = 0; j < n2 + 1; j++)
                {
                    tripletsB12.emplace_back(ii + j, ii + j, 6 * beta);
                    tripletsB12.emplace_back(ii + j, ii + j + n2 + 1, -4 * beta);
                    tripletsB12.emplace_back(ii + j, ii + j - (n2 + 1), -4 * beta);
                    tripletsB12.emplace_back(ii + j, ii + j + (n2 + 1) * 2, beta);
                    tripletsB12.emplace_back(ii + j, ii + j - (n2 + 1) * 2, beta);
                }
            }
        }

        // 处理前两行与最后两行
        for (std::size_t i = 0; i < n2 + 1; i++)
        {
            // 第一行
            tripletsB12.emplace_back(i, i, alpha + beta);
            tripletsB12.emplace_back(i, i + n2 + 1, -alpha - 2 * beta);
            tripletsB12.emplace_back(i, i + (n2 + 1) * 2, beta);

            // 第二行
            tripletsB12.emplace_back(i + n2 + 1, i, -2 * beta);
            tripletsB12.emplace_back(i + n2 + 1, i + n2 + 1, 5 * beta);
            tripletsB12.emplace_back(i + n2 + 1, i + (n2 + 1) * 2, -4 * beta);
            tripletsB12.emplace_back(i + n2 + 1, i + (n2 + 1) * 3, beta);

            // 最后一行
            auto index = n1 * (n2 + 1) + i;
            tripletsB12.emplace_back(index, index, alpha + beta);
            tripletsB12.emplace_back(index, index - (n2 + 1), -alpha - 2 * beta);
            tripletsB12.emplace_back(index, index - (n2 + 1) * 2, beta);

            // 倒数第二行
            tripletsB12.emplace_back(index - n2 - 1, index, -2 * beta);
            tripletsB12.emplace_back(index - n2 - 1, index - n2 - 1, 5 * beta);
            tripletsB12.emplace_back(index - n2 - 1, index - (n2 + 1) * 2, -4 * beta);
            tripletsB12.emplace_back(index - n2 - 1, index - (n2 + 1) * 3, beta);
        }

        auto n = (n1 + 1) * (n2 + 1);
        Eigen::SparseMatrix<double> A(m + 1, n);
        A.setFromTriplets(tripletsA.begin(), tripletsA.end());

        Eigen::SparseMatrix<double> B(n, n);
        B.setFromTriplets(tripletsB12.begin(), tripletsB12.end());

        // 构造最小二乘+能量约束系统的右端矩阵
        Eigen::MatrixXd Q(m + 1, N);
        for (std::size_t i = 0; i < m + 1; i++)
            for (std::size_t j = 0; j < N; j++)
                Q(i, j) = data[i][j];
        Q = A.transpose() * Q * gamma;

        Eigen::SparseMatrix<double> C = B + A.transpose() * A * gamma;
        Eigen::SimplicialLDLT<Eigen::SparseMatrix<double>> solver;
        C.makeCompressed();
        solver.compute(C);
        auto P = solver.solve(Q);

        // 构造 B-spline 曲面
        std::vector<Vec> cpts((n1 + 1) * (n2 + 1));
        for (std::size_t i = 0; i < n1 + 1; i++)
        {
            for (std::size_t j = 0; j < n2 + 1; j++)
            {
                auto row = i * (n2 + 1) + j;
                for (std::size_t k = 0; k < N; k++)
                    cpts[row][k] = P(row, k);
            }
        }
        m_surface = std::make_shared<Geom_BSplineSurface<N>>(std::move(cpts), std::move(uKnots), std::move(vKnots),
                                                             uDegree, vDegree);

        for (std::size_t i = 0; i < data.size(); i++)
            params[i] = Alg_Projection<N>::nearest(m_surface, data[i], params[i][0], params[i][1], Numeric::distance);

        uPoleNum *= 2;
        vPoleNum *= 2;

        wStretch *= 0.1;
        wBending *= 0.1;

        return m_surface;
    }
};

// MiniBSplineSurface *MiniCloud::SDEInitialDBS(int Nu, int Nv, int uDegree, int vDegree, Real tol, Real ru, Real rv)
// const
// {
//     // 计算一个包围盒
//     MiniArray<Vector3f> box = OrientBoundingBox(m_points);

//     // N 个控制点，m + 1 个数据点和对应参数
//     int n1 = Nu - 1;
//     int n2 = Nv - 1;
//     int m = m_points.size() - 1;

//     // 构造样条
//     MiniBSplineSurface *spline = new MiniBSplineSurface(n1 + 1, n2 + 1, uDegree, vDegree);
//     MiniArray2D<Vector3f> &P = spline->GetPoints();
//     const MiniArray<Vector3f> &Q = GetPoints();

//     // 稍微放大包围盒，不要放大原点位置
//     box[0] = box[0] * ru;
//     box[1] = box[1] * ru;
//     box[2] = box[2] * rv;
//     box[3] = box[3] * rv;

//     // 根据包围盒设置初始控制点
//     MiniArray<Vector3f> oxy(3);
//     oxy[0] = box.back() + box[1] + box[3] + box[5];
//     oxy[1] = normalize(box[0]);
//     oxy[2] = normalize(box[2]);
//     Real su = (length(box[0]) + length(box[1])) / n1;
//     Real sv = (length(box[2]) + length(box[3])) / n2;

//     for (int i = 0; i < n1 + 1; i++)
//         for (int j = 0; j < n2 + 1; j++)
//             P[i][j] = oxy[0] + oxy[1] * su * i + oxy[2] * sv * j;

//     // 设置平均权重
//     MiniArray<Real> w(m + 1);
//     for (int i = 0; i < m + 1; i++)
//         w[i] = 0.5;

//     // 加细分割曲面网格
//     n1 *= 2, n2 *= 2;

//     int times = 0;
//     Real e = 0, e1 = tol * 10;
//     while (fabs(e - e1) > tol && times < ITERATION)
//     {
//         // 记录新的误差
//         e = e1;

//         // 构造相交网格及其参数
//         MiniArray<Real> u((n1 + 1) * (n2 + 1));
//         MiniArray<Real> v((n1 + 1) * (n2 + 1));
//         MiniArray2D<Vector3f> R(n1 + 1, MiniArray<Vector3f>(n2 + 1));
//         MiniArray2D<Vector3f> N(n1 + 1, MiniArray<Vector3f>(n2 + 1));
//         for (int i = 1; i < n1; i++)
//         {
//             for (int j = 1; j < n2; j++)
//             {
//                 u[i * (n2 + 1) + j] = 1.0 / n1 * i;
//                 v[i * (n2 + 1) + j] = 1.0 / n2 * j;

//                 // 内部点计算法向
//                 MiniArray2D<Vector3f> S = spline->SurfaceDerivPoint(1, 1.0 / n1 * i, 1.0 / n2 * j);
//                 R[i][j] = S[0][0];
//                 N[i][j] = normalize(cross(S[1][0], S[0][1]));
//             }
//         }

//         // 外插计算边界点处的法向
//         // 左右边界外插
//         for (int i = 1; i < n1; i++)
//         {
//             Real t = 1.0 / n1 * i;
//             u[i * (n2 + 1)] = t;
//             v[i * (n2 + 1)] = 0;
//             R[i][0] = spline->SurfacePoint(t, 0);
//             N[i][0] = normalize(N[i][1] * 2 - N[i][2]);

//             u[i * (n2 + 1) + n2] = t;
//             v[i * (n2 + 1) + n2] = 1;
//             R[i][n2] = spline->SurfacePoint(t, 1);
//             N[i][n2] = normalize(N[i][n2-1] * 2 - N[i][n2-2]);
//         }

//         // 上下边界外插
//         for (int i = 1; i < n2; i++)
//         {
//             Real t = 1.0 / n2 * i;
//             u[i] = 0;
//             v[i] = t;
//             R[0][i] = spline->SurfacePoint(0, t);
//             N[0][i] = normalize(N[1][i] * 2 - N[2][i]);

//             u[n1 * (n2 + 1) + i] = 1;
//             v[n1 * (n2 + 1) + i] = t;
//             R[n1][i] = spline->SurfacePoint(1, t);
//             N[n1][i] = normalize(N[n1-1][i] * 2 - N[n1-2][i]);
//         }

//         // 四个角点外插
//         u[0] = 0, v[0] = 0;
//         u[n2] = 0, v[n2] = 1;
//         u[n1 * (n2 + 1)] = 1, v[n1 * (n2 + 1)] = 0;
//         u[n1 * (n2 + 1) + n2] = 1, v[n1 * (n2 + 1) + n2] = 1;

//         R[0][0] = spline->SurfacePoint(0, 0);
//         R[0][n2] = spline->SurfacePoint(0, 1);
//         R[n1][0] = spline->SurfacePoint(1, 0);
//         R[n1][n2] = spline->SurfacePoint(1, 1);

//         N[0][0] = normalize(N[0][1] * 2 - N[0][2]);
//         N[0][n2] = normalize(N[0][n2-1] * 2 - N[0][n2-2]);
//         N[n1][0] = normalize(N[n1-1][0] * 2 - N[n1-2][0]);
//         N[n1][n2] = normalize(N[n1-1][n2] * 2 - N[n1-2][n2]);

//         // 遍历每个控制点进行投影，计算新的误差
//         e1 = 0;
//         for (int i = 0; i < n1 + 1; i++)
//         {
//             for (int j = 0; j < n2 + 1; j++)
//             {
//                 Vector c(5);

//                 // 对每个点云中的点计算误差
//                 for (int mu = 0; mu < m + 1; mu++)
//                 {
//                     Real a = 1 / pow(length(R[i][j] - Q[mu]), 4);
//                     c[0] += a;
//                     c[1] += a * Q[mu].x();
//                     c[2] += a * Q[mu].y();
//                     c[3] += a * Q[mu].z();
//                     c[4] += a * (Q[mu].x() * Q[mu].x() + Q[mu].y() * Q[mu].y() + Q[mu].z() * Q[mu].z());
//                 }

//                 // 计算出极小值参数，获得近似后的点
//                 Real lambda = (c[1] * N[i][j].x() + c[2] * N[i][j].y() + c[3] * N[i][j].z()) / c[0];
//                 Real t = lambda - dot(R[i][j], N[i][j]);

//                 // 计算投影点以及误差
//                 R[i][j] = R[i][j] + N[i][j] * t;
//                 e1 += length(spline->SurfacePoint(1.0 / n1 * i, 1.0 / n2 * j) - R[i][j]);
//             }
//         }
//         e1 /= (n1 + 1) * (n2 + 1);

//         MiniCloud cloud(MiniArray<Vector3f>(0));
//         MiniArray<Vector3f> &final = cloud.GetPoints();
//         for (int i = 0; i < n1 + 1; i++)
//             for (int j = 0; j < n2 + 1; j++)
//                 final.push_back(R[i][j]);

//         // 计算完成后即可移除原先曲面
//         delete spline;

//         // 重新拟合曲面，不需要光滑化，权重为零
//         spline = cloud.ChordFittingSurface(u, v, Nu, Nv, uDegree, vDegree, w, 0, 0);
//         times++;
//     }
//     return spline;
// }

} // namespace alg

} // namespace xi
